package nsfs

import (
	"GLibHac/fssystem"
	"GLibHac/tools"
	"encoding/binary"
	"fmt"
	"github.com/pkg/errors"
)

type EntrySetInfo struct {
	Index int32
	Count int32
	End   int64
	Start int64
}

func NewEntrySetInfo(data []byte) *EntrySetInfo {
	return &EntrySetInfo{
		Index: tools.BytesToIntBigEndian(data[0:4]),
		Count: tools.BytesToIntBigEndian(data[4:8]),
		End:   int64(binary.BigEndian.Uint64(data[8:16])),
		Start: int64(binary.BigEndian.Uint64(data[16:24])),
	}
}

type EntrySetHeader struct {
	Header *NodeHeader
	Info   *EntrySetInfo
}

func NewEntrySetHeader(data []byte) *EntrySetHeader {
	return &EntrySetHeader{
		Header: NewNodeHeader(data),
		Info:   NewEntrySetInfo(data),
	}
}

const EntrySetHeaderSize = 0x18

type Offset struct {
	offset int64
	stride int
}

func (o *Offset) Plus(right int32) *Offset {
	return NewOffset(o.offset+int64(right)*int64(o.stride), o.stride)
}

func (o *Offset) Get() int64 {
	return o.offset
}

func (o *Offset) Reduction(right *Offset) int {
	return int((o.offset - right.offset) / int64(o.stride))
}

func NewOffset(offset int64, stride int) *Offset {
	return &Offset{offset, stride}
}

type StorageNode struct {
	start *Offset
	count int32
	index int32
}

func (n *StorageNode) Find(storage fssystem.IStorage, virtualAddress int64) error {
	end := n.count
	pos := n.start

	for end > 0 {
		half := end / 2
		mid := pos.Plus(half)
		offset := int64(0)
		err := storage.SetPosition(mid.Get())
		if err != nil {
			return err
		}
		offset, err = storage.ReadInt64()
		if err != nil {
			return err
		}
		if offset <= virtualAddress {
			pos = mid.Plus(1)
			end -= half + 1
		} else {
			end = half
		}
	}
	n.index = (int32)(pos.Reduction(n.start)) - 1
	return nil
}

func NewStorageNode(offset, size int64, count int32) *StorageNode {
	return &StorageNode{
		start: NewOffset(offset+NodeHeaderSize, int(size)),
		count: count,
		index: -1,
	}
}

type BucketTreeVisitor struct {
	tree *BucketTree
	*Entry
	entry         []byte
	treeOffsets   *Offsets
	entryIndex    int32
	entrySetCount int32
	entrySet      *EntrySetHeader
}

func (v *BucketTreeVisitor) Find(virtualAddress int64) (err error) {
	nodeL1 := v.tree.nodeL1.Node
	if virtualAddress >= nodeL1.EndOffset {
		return fmt.Errorf("out of range: %d", virtualAddress)
	}

	var entrySetIndex int32
	if v.tree.IsExistOffsetL2OnL1() && virtualAddress < nodeL1.BeginOffset {
		// The portion of the L2 offsets containing our target offset is stored in the L1 node
		offsets := tools.SliceLong(nodeL1.Array, nodeL1.Count)

		// Find the index of the entry containing the requested offset.
		// If the value is not found, BinarySearch will return the bitwise complement of the
		// index of the first element that is larger than the value.
		// The offsets are the start offsets of each entry, so subtracting 1 from the index of
		// the next-largest value will get us the index of the entry containing the offset.
		index := tools.BinarySearchLong(offsets, virtualAddress)
		if index < 0 {
			index = (^index) - 1
		}

		// If the requested offset comes before the first offset in the list, "index" will be -1.
		if index < 0 {
			return errors.Errorf("out of range: %d", index)
		}

		entrySetIndex = index
	} else {
		offsets := tools.SliceLong(nodeL1.Array, nodeL1.Count)
		index := tools.BinarySearchLong(offsets, virtualAddress)
		if index < 0 {
			index = (^index) - 1
		}

		if index < 0 {
			return errors.Errorf("out of range: %d", index)
		}

		if v.tree.IsExistL2() {
			if index >= v.tree.offsetCount {
				return errors.Errorf("invalid bucket tree node offset: %d", index)
			}
			entrySetIndex, err = v.FindEntrySet(virtualAddress, index)
			if err != nil {
				return err
			}
		} else {
			entrySetIndex = index
		}
	}

	// Validate the entry set index.
	if entrySetIndex < 0 || entrySetIndex >= v.tree.entrySetCount {
		return errors.Errorf("invalid bucket tree node offset: %d", entrySetIndex)
	}

	// Find the entry.
	if err = v.FindEntry(virtualAddress, entrySetIndex); err != nil {
		return err
	} else {
		v.Entry = NewEntry(v.entry)
	}
	// Set count.
	v.entrySetCount = v.tree.entrySetCount
	return nil
}

func (v *BucketTreeVisitor) FindEntrySet(virtualAddress int64, nodeIndex int32) (int32, error) {
	// Calculate node extents.
	nodeSize := v.tree.nodeSize
	nodeOffset := int64(nodeSize * (nodeIndex + 1))
	storage := v.tree.nodeStorage

	// Read and validate the header.
	if err := storage.SetPosition(nodeOffset); err != nil {
		return 0, err
	}
	bytes := make([]byte, NodeHeaderSize)
	if _, err := storage.Read(nodeOffset, bytes); err != nil {
		return 0, err
	}
	header := NewNodeHeader(bytes)
	if err := header.Verify(nodeIndex, int64(nodeSize), 8); err != nil {
		return 0, err
	}
	// Create the node, and find.
	node := NewStorageNode(nodeOffset, 8, header.EntryCount)
	if err := node.Find(storage, virtualAddress); err != nil {
		return 0, err
	}

	if node.index < 0 {
		return 0, errors.Errorf("invalid bucket tree virtual offset: %d", node.index)
	}

	// Return the index.
	return v.tree.GetEntrySetIndex(header.Index, node.index), nil
}

func (v *BucketTreeVisitor) FindEntry(virtualAddress int64, entrySetIndex int32) error {
	// Calculate entry set extents.
	entrySize := v.tree.entrySize
	entrySetSize := v.tree.nodeSize
	entrySetOffset := entrySetSize * entrySetIndex
	storage := v.tree.entryStorage

	// Read and validate the entry set.
	bytes := make([]byte, NodeHeaderSize)
	_, err := storage.Read(int64(entrySetOffset), bytes)
	if err != nil {
		return err
	}
	entrySet := NewEntrySetHeader(bytes)
	err = entrySet.Header.Verify(entrySetIndex, int64(entrySetSize), int64(entrySize))
	if err != nil {
		return err
	}

	// Create the node, and find.
	var node = NewStorageNode(int64(entrySetOffset), int64(entrySize), entrySet.Info.Count)
	err = node.Find(storage, virtualAddress)
	if err != nil {
		return err
	}

	if node.index < 0 {
		return errors.Errorf("invalid bucket tree virtual offset: %d", node.index)
	}

	// Copy the data into entry.
	v.entryIndex = -1
	entryIndex := node.index
	var entryOffset = v.GetBucketTreeEntryOffset(entrySetOffset, entrySize, entryIndex)
	if _, err := storage.Read(entryOffset, v.entry); err != nil {
		return err
	}

	// Set our entry set/index.
	v.entrySet = entrySet
	v.entryIndex = entryIndex
	return nil
}

func (v *BucketTreeVisitor) GetBucketTreeEntryOffset(entrySetOffset, entrySize, entryIndex int32) int64 {
	return int64(entrySetOffset) + int64(NodeHeaderSize) + int64(entryIndex)*int64(entrySize)
}

func (v *BucketTreeVisitor) CanMoveNext() bool {
	return v.IsValid() && (v.entryIndex+1 < v.entrySet.Info.Count || v.entrySet.Info.Index+1 < v.entrySetCount)
}

func (v *BucketTreeVisitor) IsValid() bool {
	return v.entryIndex >= 0
}

func (v *BucketTreeVisitor) MoveNext() error {
	if !v.IsValid() {
		return errors.Errorf("out of range")
	}

	entryIndex := v.entryIndex + 1

	// Invalidate our index, and read the header for the next index.
	if entryIndex == v.entrySet.Info.Count {
		entrySetIndex := v.entrySet.Info.Index + 1
		if entrySetIndex >= v.entrySetCount {
			return errors.Errorf("out of range: %d", entryIndex)
		}
		v.entryIndex = -1

		end := v.entrySet.Info.End

		entrySetSize := v.tree.nodeSize
		entrySetOffset := entrySetIndex * entrySetSize

		bytes := make([]byte, EntrySetHeaderSize)
		if _, err := v.tree.entryStorage.Read(int64(entrySetOffset), bytes); err != nil {
			return err
		}
		v.entrySet = NewEntrySetHeader(bytes)
		if err := v.entrySet.Header.Verify(entrySetIndex, int64(entrySetSize), int64(v.tree.entrySize)); err != nil {
			return err
		}
		if v.entrySet.Info.Start != end || v.entrySet.Info.Start >= v.entrySet.Info.End {
			return errors.Errorf("invalid bucket tree entry set offset: %d", v.entrySet.Info.Start)
		}
		entryIndex = 0
	} else {
		v.entryIndex = 1
	}

	// Read the new entry
	entrySize := v.tree.entrySize
	entryOffset := v.GetBucketTreeEntryOffset(v.entrySet.Info.Index*v.tree.nodeSize, entrySize, entryIndex)
	bytes := make([]byte, EntrySetHeaderSize)
	if _, err := v.tree.entryStorage.Read(entryOffset, bytes); err != nil {
		return err
	}
	v.entrySet = NewEntrySetHeader(bytes)
	// Note that we changed index.
	v.entryIndex = entryIndex
	return nil
}

func NewBucketTreeVisitor(tree *BucketTree, offsets *Offsets) *BucketTreeVisitor {
	return &BucketTreeVisitor{
		tree:        tree,
		treeOffsets: offsets,
		entry:       make([]byte, tree.entrySize),
	}
}

const BucketTreeHeaderMagic = "BKTR"
const BucketTreeHeaderMaxVersion = 1

type BucketTreeHeader struct {
	Magic      string
	Version    uint32
	EntryCount int32
}

func (h *BucketTreeHeader) Verify() error {
	if h.Magic != BucketTreeHeaderMagic {
		return errors.Errorf("invalid Bucket Tree Signatur: %s", h.Magic)
	}

	if h.EntryCount < 0 {
		return errors.Errorf("invalid Bucket Tree Entry Count: %d", h.EntryCount)
	}
	if h.Version > BucketTreeHeaderMaxVersion {
		return errors.Errorf("invalid Bucket Tree Version: %d", h.Version)
	}
	return nil
}

func NewBucketTreeHeader(data []byte) *BucketTreeHeader {
	return &BucketTreeHeader{
		Magic:      string(data[0:4]),
		Version:    binary.LittleEndian.Uint32(data[4:]),
		EntryCount: tools.BytesToIntLittleEndian(data[8:]),
	}
}

const NodeHeaderSize = 0x10

type NodeHeader struct {
	Index      int32
	EntryCount int32
	OffsetEnd  int64
}

func (n *NodeHeader) Verify(nodeIndex int32, nodeSize, entrySize int64) error {
	if n.Index != nodeIndex {
		return errors.Errorf("InvalidBucketTreeNodeIndex, excepted %d, but got %d", nodeIndex, n.Index)
	}

	if entrySize == 0 || nodeSize < entrySize+NodeHeaderSize {
		return errors.Errorf("invalid size:  %d", entrySize)
	}

	maxEntryCount := (nodeSize - NodeHeaderSize) / entrySize

	if n.EntryCount <= 0 || maxEntryCount < int64(n.EntryCount) {
		return errors.Errorf("invalid bucket tree node entry count: %d", n.EntryCount)
	}

	if n.OffsetEnd < 0 {
		return errors.Errorf("invalid bucket tree node offset: %d", n.OffsetEnd)
	}
	return nil
}

func NewNodeHeader(data []byte) *NodeHeader {
	return &NodeHeader{
		Index:      tools.BytesToIntLittleEndian(data[0:4]),
		EntryCount: tools.BytesToIntLittleEndian(data[4:8]),
		OffsetEnd:  int64(binary.LittleEndian.Uint64(data[8:])),
	}
}

type BucketTreeNode struct {
	Count         int32
	L2BeginOffset int64
	BeginOffset   int64
	EndOffset     int64
	Array         []byte
}

func NewBucketTreeNode(buffer *NodeBuffer) (r *BucketTreeNode) {
	r = &BucketTreeNode{
		Array:     buffer.buffer,
		EndOffset: buffer.Header.OffsetEnd,
		Count:     buffer.Header.EntryCount,
	}
	r.L2BeginOffset = int64(binary.LittleEndian.Uint64(buffer.buffer[r.Count*8:]))
	r.BeginOffset = int64(binary.LittleEndian.Uint64(buffer.buffer))
	return
}

type NodeBuffer struct {
	buffer []byte
	Header *NodeHeader
	Node   *BucketTreeNode
}

func NewNodeBuffer(buffer []byte) (r *NodeBuffer) {
	r = &NodeBuffer{
		buffer: buffer,
		Header: NewNodeHeader(buffer),
	}
	r.Node = NewBucketTreeNode(r)
	return
}

type Offsets struct {
	StartOffset int64
	EndOffset   int64
}

func NewOffsets(start, end int64) (r *Offsets) {
	return &Offsets{start, end}
}

func (o *Offsets) IsIncludeOffset(offset int64) bool {
	return o.StartOffset <= offset && offset < o.EndOffset
}

func (o *Offsets) IsInclude(offset, size int64) bool {
	return size > 0 && o.StartOffset <= offset && size <= o.EndOffset-offset
}

type OffsetCache struct {
	*Offsets
	IsInitialized bool
}

func NewOffsetCache() (r *OffsetCache) {
	return &OffsetCache{
		Offsets: NewOffsets(-1, -1),
	}
}

type BucketTree struct {
	nodeL1        *NodeBuffer
	offsetCache   *OffsetCache
	nodeStorage   fssystem.IStorage
	entryStorage  fssystem.IStorage
	nodeSize      int32
	entrySize     int32
	entryCount    int32
	offsetCount   int32
	entrySetCount int32
}

func (t *BucketTree) GetOffsetCount(nodeSize int32) int32 {
	return (nodeSize - NodeHeaderSize) / 8
}

func (t *BucketTree) GetOffsets() *Offsets {
	return t.offsetCache.Offsets
}

func (t *BucketTree) IsExistOffsetL2OnL1() bool {
	return t.IsExistL2() && t.nodeL1.Header.EntryCount < t.offsetCount
}

func (t *BucketTree) IsExistL2() bool {
	return t.offsetCount < t.entrySetCount
}

func (t *BucketTree) GetEntrySetIndex(nodeIndex int32, offsetIndex int32) int32 {
	return (t.offsetCount - t.nodeL1.Header.EntryCount) + (t.offsetCount * nodeIndex) + offsetIndex
}

func GetEntrySetCount(nodeSize, entrySize, entryCount int32) int32 {
	var entryCountPerNode = GetEntryCount(nodeSize, entrySize)
	return int32(tools.DivideUp(int64(entryCount), int64(entryCountPerNode)))
}

func GetEntryCount(nodeSize, entrySize int32) int32 {
	return (nodeSize - NodeHeaderSize) / entrySize
}

func NewBucketTree(nodeStorage, entryStorage fssystem.IStorage, nodeSize, entrySize, entryCount int32) (*BucketTree, error) {
	if entryCount <= 0 {
		return nil, errors.Errorf("entryCount must be greater than zero")
	}
	buffer := make([]byte, nodeSize)
	if _, err := nodeStorage.Read(0, buffer); err != nil {
		return nil, err
	}
	r := &BucketTree{
		nodeStorage:  nodeStorage,
		entryStorage: entryStorage,
		offsetCache:  NewOffsetCache(),
	}
	r.nodeL1 = NewNodeBuffer(buffer)
	if err := r.nodeL1.Header.Verify(0, int64(nodeSize), 8); err != nil {
		return nil, err
	}
	// Validate offsets.
	var offsetCount = r.GetOffsetCount(nodeSize)
	var entrySetCount = GetEntrySetCount(nodeSize, entrySize, entryCount)
	node := r.nodeL1.Node
	var startOffset int64
	if offsetCount < entrySetCount && node.Count < offsetCount {
		startOffset = node.L2BeginOffset
	} else {
		startOffset = node.BeginOffset
	}
	endOffset := node.EndOffset

	if startOffset < 0 || startOffset > node.BeginOffset || startOffset >= endOffset {
		return nil, errors.Errorf("invalid bucket tree entry offset %d ", startOffset)
	}
	r.nodeStorage = nodeStorage
	r.entryStorage = entryStorage
	r.nodeSize = nodeSize
	r.entrySize = entrySize
	r.entryCount = entryCount
	r.offsetCount = offsetCount
	r.entrySetCount = entrySetCount
	r.offsetCache.IsInitialized = true
	r.offsetCache.Offsets.StartOffset = startOffset
	r.offsetCache.Offsets.EndOffset = endOffset
	return r, nil
}
